module Main where

import SliceADT
import Utils
import SliceSummary
import InterIRSlice

import System.Environment (getEnv,getArgs)
import Options.Applicative
import Data.Map ( Map )
import qualified Data.Map as M
import Data.IntMap ( IntMap )
import qualified Data.IntMap as IM
import Data.IntSet ( IntSet )
import qualified Data.IntSet as IS
import Data.Maybe  
import Data.List 
import Data.Char ( isLetter )

import LLVM.Analysis
import LLVM.Analysis.Util.Testing
import LLVM.Parse

 
-----
---
data Direction = Fwd | Bwd | Both deriving (Read, Ord, Eq, Show)

data Opts = Opts { outputFile :: Maybe FilePath
                 , sliceVars :: [String]
                 , direction :: Direction
                 , inputFile :: FilePath
                 } deriving (Show)

cmdOpts :: Parser Opts
cmdOpts = Opts
 <$> optional (strOption
     ( long "output"
     <> short 'o'
     <> metavar "FILE/DIR"
     <> help "The destination of a file output"))
  <*>  many (strOption
      ( long "criterion"
      <> short 'c'
      <> metavar "VARIABLE"
      <> help "The criterion variables (with the style of Var@Fun,e.g. x@main) for slicing. (If null, just output the slice table for all single variables.)"))
  <*> option
      ( long "direction"
      <> short 'd'
      <> metavar "DIRECTION"
      <> help "The type of output to slicing: Fwd, Bwd, Both. Default: Bwd"
      <> value Bwd)
  <*> argument str 
      ( metavar "FILE" 
      <> help "The input file which can be bitcode,llvm assembly, or C/CPP sourcecode")

main :: IO ()
main = execParser args >>= runMain
  where
    args = info (helper <*> cmdOpts)
      ( fullDesc
      <> progDesc "Generate static slice based on paramiterized method for FILE (which can be bitcode,llvm assembly, or C/CPP sourcecode)"
      <> header "ParaSlicer - Paramiterized program slicing for LLVM IR")

runMain :: Opts -> IO ()
runMain opts = do
    let dir = direction opts
        inFile = inputFile opts
        outFile = outputFile opts
        s = sliceVars opts      
        --    
        sortFst = sortBy (\(x,_) (y,_) -> compare x y)           
        groupFst = groupBy (\(x,_) (y,_) -> x == y)    
        mergSrc [] = ""
        mergSrc ss = hd ++ ": " ++ show lns  
         where hd = fst (head ss)
               lns = nub . sort $ concatMap snd ss   
--        inst2lines = map mergSrc. groupFst. sortFst. mapMaybe valueSrcLn. findVals valMap

    m <- file2module inFile
    let summ = modularSlice m sliceTransfer
        valMap = genValueMap m                           
        res = if null s then tblRes  else sliRes
        
        tblRes = case dir of
           Fwd -> fwdTbl
           Bwd -> bwdTbl
           Both -> bwdTbl ++ "\n" ++ fwdTbl                   
          where bwdTbl = "\nBackward Static SliceTable:" ++ showTbl (_bwdSliceTable summ)
                fwdTbl = "\nForward Static SliceTable:" ++ showTbl (_fwdSliceTable summ)
        showTbl = showSrcLns.   -- showLns  -- for srcLine
--                  show . avgSize .     -- for AVG slice size
                  M.mapKeys (drop 1). M.filterWithKey (\(_:k:_) _->(isLetter k || k == '_'))      
          where avgSize m = round $ (allSlices m)/(fromIntegral $ M.size m)  
                allSlices = fromIntegral . sum . map IS.size . M.elems                 
--                showLns = show . M.map (IS.fromList. concatMap valueLine. findVals valMap) 
                showSrcLns = show . M.map (map mergSrc. groupFst. sortFst.
                                 mapMaybe valueSrcLn. findVals valMap)  
                                 
        sliRes = case dir of
           Fwd -> fwdSli
           Bwd -> bwdSli
           Both -> bwdSli ++ "\n" ++ fwdSli                   
          where bwdSli = "Backward Static Slice for " ++ show s ++ showSli (_bwdSliceTable summ)
                fwdSli = "Forward Static Slice for " ++ show s ++ showSli (_fwdSliceTable summ)
        showSli slices = ":\n" ++ srcLines ++ valStrings  --  show vals
          where vals = findVals valMap. unionLkpSli s . M.mapKeys (drop 1) $ slices
                srcLines = " <SourceLines> " ++ show (map mergSrc. groupFst. sortFst.
                                 mapMaybe valueSrcLn $ vals) ++ ": \n"
--                srcLines = " <SourceLines> " ++ show (nub. concatMap valueLine $ vals) ++ ": \n" 
                valStrings = concatMap showVal vals
                showVal v = "  " ++ show v ++ "\n"   
                 
    case outFile of 
      Nothing -> putStrLn res
      Just ofile -> appendFile ofile res
--    printInfo m


file2module :: FilePath -> IO Module    
file2module fn = buildModule ["-g"] optOptions fn2mod fn
   where optOptions = ["-gvn", "-basicaa"]     -- ["-mem2reg", "-gvn"] 
         fn2mod = parseLLVMFile defaultParserOptions


